function [bci,trainDat]=bciGetTrainDat(bci,varargin)

% [bci,trainDat]=bciGetTrainDat(bci,varargin)
% PURPOSE:
% Transform epoched data classifier conform data, merge runs deselect bad epochs
% and bad channels, do balance trainset, select features
% INPUT:
% bci:      Structure holding parameters and variables
% varargin is expected to be a cell array with one element for each run
% the data structure of a run is created e.g. in bciGetTrainBuffer and
% saved to disk if selected to do so


nEpochs=[];
if length(varargin)==1 && iscell(varargin{1}),
    varargin=varargin{1};
end

for k=1:length(varargin),
    nEpochs(k)=size(varargin{k}.trainEpoch,3);
    if isfield(varargin{k},'isRetrainTrial')&&sum(varargin{k}.isRetrainTrial)>0, % comes from feedback session
        varargin{k}.trainLabel(~varargin{k}.isRetrainTrial)=0;
    end
end

eventsToClassify = bci.eventsToClassify;

badChanMap = zeros(size(varargin{k}.badChanMap,1),sum(nEpochs));
badChan=[];
label = zeros(1,sum(nEpochs));
selector = zeros(1,sum(nEpochs));
epochIdxList = zeros(1,sum(nEpochs));

% merge labels of different runs
for k=1:length(varargin),
    label(sum(nEpochs(1:k-1))+1:sum(nEpochs(1:k)))=varargin{k}.trainLabel;
    selector(sum(nEpochs(1:k-1))+1:sum(nEpochs(1:k)))=k;
    epochIdxList(sum(nEpochs(1:k-1))+1:sum(nEpochs(1:k)))=1:nEpochs(k);
    badChanMap(:,sum(nEpochs(1:k-1))+1:sum(nEpochs(1:k)))=varargin{k}.badChanMap;
    badChan = union(badChan,find(mean(varargin{k}.badChanMap,2)>bci.param.badChanThresh));
    % find dead channels
    sumThroughTrials = sum(abs(varargin{k}.trainEpoch),3);
    deadChan = sum(sumThroughTrials==0,2);
    if sum(deadChan)>0,
        fprintf('No signal in channels:\n');
        disp(find(deadChan));
        badChan=union(badChan,find(deadChan));
    end
end
%bci.goodChan=setdiff(bci.chanOfInterest,badChan); % bug: wrong mapping ?
% bci.goodChan=setdiff(1:size(badChanMap,1),badChan);
[dummy,posInMap]=intersect(bci.preprocChan,bci.chanOfInterest);
bci.goodChan=intersect(setdiff(1:size(badChanMap,1),badChan),posInMap);

goodEpoch=find(sum(badChanMap(bci.goodChan,:),1)==0);

formatStr = repmat(' %i',1,length(eventsToClassify));
sizePerClass = sum((eventsToClassify'*ones(1,length(goodEpoch))-...
                   (ones(length(eventsToClassify),1)*label(goodEpoch)))==0,2);
fprintf(['badEpochs: %i, bad channels: %i\n',...
         'goodEpochs cond(' formatStr '):' formatStr '\n'],...
         sum(nEpochs)-length(goodEpoch),length(badChan),eventsToClassify, sizePerClass);

if bci.param.bootstrapSize>0,
    bootstrapIdx = getBootstrappedTrainset(label(goodEpoch),bci.param.bootstrapSize);
    goodEpoch = [goodEpoch, goodEpoch(bootstrapIdx)];
end

if bci.param.balanceTrainSet,
    % remove older samples
    label_bin = getBalancedTrainset(label(goodEpoch),bci.param.maxTrainSamp,'leading');%'equal');
    goodEpoch(~label_bin)=[];
end

label = label(goodEpoch);
selector = selector(goodEpoch);

% merge labels of different epochs
epoch = zeros(length(bci.goodChan),length(bci.init.featOfInterest),length(goodEpoch));
for k=1:length(varargin),
    curEpIdx = epochIdxList(goodEpoch(selector==k));
    newEpIdx = selector==k;
    epoch(:,:,newEpIdx)=varargin{k}.trainEpoch(bci.goodChan,bci.init.featOfInterest,curEpIdx);
end

clear varargin

fprintf(['Trainset size cond(' formatStr '):' formatStr '\n'],...
        eventsToClassify, sum((eventsToClassify'*ones(1,length(label))-...
                                 (ones(length(eventsToClassify),1)*label))==0,2));

% set labels
trainDat.label = label';
trainDat.selector = selector';
% do feature selection
bci.selectedFeat = intersect(1:bci.init.nFeatPerChan,bci.init.featOfInterest);
if ~isempty(bci.param.featureSelection),
    if ~iscell(bci.param.featureSelection),
        bci.param.featureSelection={bci.param.featureSelection};
    end
    for featSelK=1:length(bci.param.featureSelection),
        featSelStr = bci.param.featureSelection{featSelK};
        chanStrPos=strfind(featSelStr,'chan');
        freqStrPos=strfind(featSelStr,'freq');
        if isempty(freqStrPos),
            freqStrPos=strfind(featSelStr,'band');
        end
        % do channel selection
        if ~isempty(chanStrPos),
            nChans = str2double(featSelStr(chanStrPos+4:end));
            methodStr=featSelStr(1:chanStrPos-1);
            if bci.param.CSPfilter > 0, %do channel selection in frequency space
                fftEpoch = zeros(size(epoch,1),length(bci.init.fftFreqIdx),size(epoch,3));
                for tr = 1:size(epoch,3),
                    % Fourier transform
                    Y = fft(epoch(:,:,tr).*(ones(size(epoch,1),1)*bci.init.fftWin),bci.init.NFFT,2);
                    tmp=abs(Y).^2;
                    fftEpoch(:,:,tr)=tmp(:,bci.init.fftFreqIdx);
                end
                [trainDat.dat bci] = bciTransformData(fftEpoch,bci,0);
            else
                [trainDat.dat bci] = bciTransformData(epoch,bci,0);
            end
            bci.featureMask=true(1,size(trainDat.dat,2));
            selectedChan = bciFeatureRanking({bci,trainDat},nChans,methodStr,'goodchan','max');
            bci.goodChan = bci.goodChan(selectedChan);
            epoch(setdiff(1:size(epoch,1),selectedChan),:,:)=[];
        end
        % do frequency selection
        if ~isempty(freqStrPos),
            nFreqs = str2double(featSelStr(freqStrPos+4:end));
            methodStr=featSelStr(1:freqStrPos-1);
            if bci.param.CSPfilter > 0, %do channel selection in frequency space
                error('When using CSP, no frequencies can be selected. Check bci.param.featureSelection');
            else
                [trainDat.dat bci] = bciTransformData(epoch,bci,0);
            end
            bci.featureMask=true(1,size(trainDat.dat,2));
            [selectedFreq selectedFreqIdx]= bciFeatureRanking({bci,trainDat},nFreqs,methodStr,featSelStr(freqStrPos:freqStrPos+3),'max');
            bci.selectedFeat = bci.selectedFeat(selectedFreqIdx);
            epoch(:,setdiff(1:size(epoch,2),selectedFreqIdx),:)=[];
        end
    end
end

if bci.param.CSPfilter > 0, %estimate and apply csp
    if strcmpi(bci.param.classifier,'svm'),
        idx=label'==eventsToClassify(1)|label'==eventsToClassify(2);
        label=label(idx);
    else
        idx=true(size(label))';
    end
    [epoch,bci]=bciCSPtransform(bci,epoch(:,:,idx),label');
end

% get normalized data
[trainDat.dat bci] = bciTransformData(epoch,bci,0);


bci.featureMask=true(1,size(trainDat.dat,2));
% do feature selection of raw features, i.e. in feature vector space
if ~isempty(bci.param.featureSelection),
    if ~iscell(bci.param.featureSelection),
        bci.param.featureSelection={bci.param.featureSelection};
    end
    for featSelK=1:length(bci.param.featureSelection),
        featSelStr = bci.param.featureSelection{featSelK};
        rawStrPos=strfind(featSelStr,'raw');    
        if ~isempty(rawStrPos),
            nFeat = str2double(featSelStr(rawStrPos+4:end));
            methodStr=featSelStr(1:rawStrPos-1);
            % rank features
            [selectedFeat selectedFeatIdx]= bciFeatureRanking({bci,trainDat},nFeat,methodStr,'raw','max');
            if ~isempty(selectedFeatIdx),
                bci.featureMask(:) = false;
                bci.featureMask(selectedFeatIdx)=true;
                trainDat.dat=trainDat.dat(:,bci.featureMask);
            end
        end
    end
end